import logging
import os
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
from typing import Dict, List, Tuple, Callable, Union, Any, Optional, Type
from concurrent.futures import ProcessPoolExecutor, as_completed
import itertools
from tqdm import tqdm

# 导入QEngine和其他必要组件
from qtorch.core.qengine import QEngine
from qtorch.core.interfaces import IStrategy

class QOptimizer:
    """
    参数优化框架，支持网格搜索、随机搜索和贝叶斯优化等方法
    """
    
    def __init__(self, strategy_class: Type[IStrategy], symbols: List[str], 
               start_date: str, end_date: str, initial_cash: float = 1000000,
               commission: float = 0.0003, output_dir: str = "./output/optimization"):
        """
        初始化参数优化器
        
        参数:
            strategy_class: 要优化的策略类
            symbols: 交易标的列表
            start_date: 回测开始日期
            end_date: 回测结束日期
            initial_cash: 初始资金
            commission: 交易佣金
            output_dir: 输出目录
        """
        self.strategy_class = strategy_class
        self.symbols = symbols if isinstance(symbols, list) else [symbols]
        self.start_date = start_date
        self.end_date = end_date
        self.initial_cash = initial_cash
        self.commission = commission
        self.output_dir = output_dir
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 初始化评估指标
        self.metrics_history = []
        self.best_params = None
        self.best_metric = None
        self.best_result = None
        
        logging.info(f"参数优化器初始化完成: {strategy_class.__name__}, {len(symbols)}个交易标的")
    
    def run_backtest(self, params: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, float], pd.DataFrame]:
        """
        运行单次回测
        
        参数:
            params: 策略参数字典
            
        返回:
            Tuple: (策略参数, 绩效指标, 回测数据)
        """
        # 创建引擎
        engine = QEngine()
        
        # 添加数据
        for symbol in self.symbols:
            engine.add_data(symbol, self.start_date, self.end_date)
        
        # 设置参数
        engine.set_cash(self.initial_cash)
        engine.set_commission(self.commission)
        
        # 添加策略
        engine.add_strategy(self.strategy_class, **params)
        
        # 运行回测
        result = engine.run()
        
        return params, result.metrics, result.data
    
    def grid_search(self, param_grid: Dict[str, List[Any]], 
                  metric: str = "夏普比率", higher_is_better: bool = True,
                  n_jobs: int = 1, verbose: bool = True) -> Dict[str, Any]:
        """
        网格搜索优化
        
        参数:
            param_grid: 参数网格, 如 {'fast_period': [5, 10, 20], 'slow_period': [50, 100, 200]}
            metric: 优化目标指标
            higher_is_better: 是否越高越好
            n_jobs: 并行任务数
            verbose: 是否显示进度
            
        返回:
            Dict: 最佳参数
        """
        # 生成所有参数组合
        param_names = list(param_grid.keys())
        param_values = list(param_grid.values())
        param_combinations = list(itertools.product(*param_values))
        
        total_combinations = len(param_combinations)
        logging.info(f"网格搜索: 共{total_combinations}种参数组合")
        
        if verbose:
            print(f"开始网格搜索, 共{total_combinations}种参数组合...")
        
        # 清空历史
        self.metrics_history = []
        
        start_time = time.time()
        
        if n_jobs > 1:
            # 多进程并行处理
            futures = []
            with ProcessPoolExecutor(max_workers=n_jobs) as executor:
                # 提交所有任务
                for values in param_combinations:
                    params = dict(zip(param_names, values))
                    futures.append(executor.submit(self.run_backtest, params))
                
                # 收集结果
                for i, future in enumerate(tqdm(as_completed(futures), total=len(futures), disable=not verbose)):
                    try:
                        params, metrics, data = future.result()
                        self.metrics_history.append({
                            'params': params,
                            'metrics': metrics,
                            'data': data
                        })
                        
                        # 显示进度
                        if verbose and (i+1) % max(1, total_combinations//10) == 0:
                            elapsed = time.time() - start_time
                            remaining = elapsed / (i+1) * (total_combinations - i - 1)
                            print(f"进度: {i+1}/{total_combinations}, "
                                 f"耗时: {elapsed:.1f}秒, 预计剩余: {remaining:.1f}秒")
                    except Exception as e:
                        logging.error(f"参数组合评估失败: {str(e)}", exc_info=True)
        else:
            # 单进程顺序处理
            for i, values in enumerate(tqdm(param_combinations, disable=not verbose)):
                params = dict(zip(param_names, values))
                try:
                    result_params, metrics, data = self.run_backtest(params)
                    self.metrics_history.append({
                        'params': result_params,
                        'metrics': metrics,
                        'data': data
                    })
                    
                    # 显示进度
                    if verbose and (i+1) % max(1, total_combinations//10) == 0:
                        elapsed = time.time() - start_time
                        remaining = elapsed / (i+1) * (total_combinations - i - 1)
                        print(f"进度: {i+1}/{total_combinations}, "
                             f"耗时: {elapsed:.1f}秒, 预计剩余: {remaining:.1f}秒")
                except Exception as e:
                    logging.error(f"参数组合评估失败: {str(e)}", exc_info=True)
        
        # 寻找最佳参数
        if higher_is_better:
            best_idx = max(range(len(self.metrics_history)), 
                         key=lambda i: self.metrics_history[i]['metrics'].get(metric, -np.inf))
        else:
            best_idx = min(range(len(self.metrics_history)), 
                         key=lambda i: self.metrics_history[i]['metrics'].get(metric, np.inf))
        
        self.best_params = self.metrics_history[best_idx]['params']
        self.best_metric = self.metrics_history[best_idx]['metrics'].get(metric)
        self.best_result = self.metrics_history[best_idx]
        
        elapsed = time.time() - start_time
        
        if verbose:
            print(f"\n网格搜索完成, 耗时: {elapsed:.1f}秒")
            print(f"最佳参数: {self.best_params}")
            print(f"最佳{metric}: {self.best_metric:.4f}")
        
        logging.info(f"网格搜索完成, 耗时: {elapsed:.1f}秒")
        logging.info(f"最佳参数: {self.best_params}")
        logging.info(f"最佳{metric}: {self.best_metric:.4f}")
        
        # 生成优化报告
        self.generate_optimization_report(metric, higher_is_better)
        
        return self.best_params
    
    def random_search(self, param_distributions: Dict[str, Callable[[], Any]], n_iter: int = 20,
                     metric: str = "夏普比率", higher_is_better: bool = True,
                     n_jobs: int = 1, verbose: bool = True) -> Dict[str, Any]:
        """
        随机搜索优化
        
        参数:
            param_distributions: 参数分布, 如 {'fast_period': lambda: np.random.randint(5, 30)}
            n_iter: 迭代次数
            metric: 优化目标指标
            higher_is_better: 是否越高越好
            n_jobs: 并行任务数
            verbose: 是否显示进度
            
        返回:
            Dict: 最佳参数
        """
        logging.info(f"随机搜索: 计划{n_iter}次迭代")
        
        if verbose:
            print(f"开始随机搜索, 计划{n_iter}次迭代...")
        
        # 清空历史
        self.metrics_history = []
        
        start_time = time.time()
        
        # 生成随机参数
        param_combinations = []
        for _ in range(n_iter):
            params = {name: dist() for name, dist in param_distributions.items()}
            param_combinations.append(params)
        
        if n_jobs > 1:
            # 多进程并行处理
            futures = []
            with ProcessPoolExecutor(max_workers=n_jobs) as executor:
                # 提交所有任务
                for params in param_combinations:
                    futures.append(executor.submit(self.run_backtest, params))
                
                # 收集结果
                for i, future in enumerate(tqdm(as_completed(futures), total=len(futures), disable=not verbose)):
                    try:
                        params, metrics, data = future.result()
                        self.metrics_history.append({
                            'params': params,
                            'metrics': metrics,
                            'data': data
                        })
                        
                        # 显示进度
                        if verbose and (i+1) % max(1, n_iter//5) == 0:
                            elapsed = time.time() - start_time
                            remaining = elapsed / (i+1) * (n_iter - i - 1)
                            print(f"进度: {i+1}/{n_iter}, "
                                 f"耗时: {elapsed:.1f}秒, 预计剩余: {remaining:.1f}秒")
                    except Exception as e:
                        logging.error(f"参数组合评估失败: {str(e)}", exc_info=True)
        else:
            # 单进程顺序处理
            for i, params in enumerate(tqdm(param_combinations, disable=not verbose)):
                try:
                    result_params, metrics, data = self.run_backtest(params)
                    self.metrics_history.append({
                        'params': result_params,
                        'metrics': metrics,
                        'data': data
                    })
                    
                    # 显示进度
                    if verbose and (i+1) % max(1, n_iter//5) == 0:
                        elapsed = time.time() - start_time
                        remaining = elapsed / (i+1) * (n_iter - i - 1)
                        print(f"进度: {i+1}/{n_iter}, "
                             f"耗时: {elapsed:.1f}秒, 预计剩余: {remaining:.1f}秒")
                except Exception as e:
                    logging.error(f"参数组合评估失败: {str(e)}", exc_info=True)
        
        # 寻找最佳参数
        if higher_is_better:
            best_idx = max(range(len(self.metrics_history)), 
                         key=lambda i: self.metrics_history[i]['metrics'].get(metric, -np.inf))
        else:
            best_idx = min(range(len(self.metrics_history)), 
                         key=lambda i: self.metrics_history[i]['metrics'].get(metric, np.inf))
        
        self.best_params = self.metrics_history[best_idx]['params']
        self.best_metric = self.metrics_history[best_idx]['metrics'].get(metric)
        self.best_result = self.metrics_history[best_idx]
        
        elapsed = time.time() - start_time
        
        if verbose:
            print(f"\n随机搜索完成, 耗时: {elapsed:.1f}秒")
            print(f"最佳参数: {self.best_params}")
            print(f"最佳{metric}: {self.best_metric:.4f}")
        
        logging.info(f"随机搜索完成, 耗时: {elapsed:.1f}秒")
        logging.info(f"最佳参数: {self.best_params}")
        logging.info(f"最佳{metric}: {self.best_metric:.4f}")
        
        # 生成优化报告
        self.generate_optimization_report(metric, higher_is_better)
        
        return self.best_params
    
    def generate_optimization_report(self, metric: str = "夏普比率", 
                                   higher_is_better: bool = True) -> str:
        """
        生成优化报告
        
        参数:
            metric: 优化目标指标
            higher_is_better: 是否越高越好
            
        返回:
            str: 报告文件路径
        """
        if not self.metrics_history:
            logging.warning("没有优化历史，无法生成报告")
            return ""
        
        # 提取参数和指标
        results = []
        for entry in self.metrics_history:
            result = {'metric': entry['metrics'].get(metric, np.nan)}
            result.update(entry['params'])
            results.append(result)
        
        # 转换为DataFrame
        df_results = pd.DataFrame(results)
        
        # 按指标排序
        if higher_is_better:
            df_results = df_results.sort_values('metric', ascending=False)
        else:
            df_results = df_results.sort_values('metric', ascending=True)
        
        # 保存结果
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_path = os.path.join(self.output_dir, f"optimization_results_{timestamp}.csv")
        df_results.to_csv(csv_path, index=False)
        
        # 生成散点矩阵图
        if len(df_results.columns) > 2:  # 至少有一个参数和指标
            try:
                # 选择前N个结果进行可视化
                top_n = min(100, len(df_results))
                plot_df = df_results.head(top_n)
                
                # 创建多个子图
                params = [col for col in plot_df.columns if col != 'metric']
                n_params = len(params)
                
                if n_params > 1:
                    # 创建散点矩阵
                    fig, axes = plt.subplots(n_params, n_params, figsize=(n_params*3, n_params*3))
                    
                    # 填充矩阵
                    for i, param_i in enumerate(params):
                        for j, param_j in enumerate(params):
                            if i == j:
                                # 对角线: 直方图
                                axes[i, j].hist(plot_df[param_i], bins=10)
                                axes[i, j].set_title(param_i)
                            else:
                                # 非对角线: 散点图
                                sc = axes[i, j].scatter(
                                    plot_df[param_j], 
                                    plot_df[param_i], 
                                    c=plot_df['metric'], 
                                    cmap='viridis', 
                                    alpha=0.7
                                )
                                axes[i, j].set_xlabel(param_j)
                                axes[i, j].set_ylabel(param_i)
                    
                    # 添加颜色条
                    cbar = fig.colorbar(sc, ax=axes.ravel().tolist())
                    cbar.set_label(metric)
                    
                    plt.tight_layout()
                    plot_path = os.path.join(self.output_dir, f"optimization_scatter_matrix_{timestamp}.png")
                    plt.savefig(plot_path, dpi=150)
                    plt.close()
                
                # 对每个参数创建与指标的关系图
                for param in params:
                    fig, ax = plt.subplots(figsize=(10, 6))
                    ax.scatter(plot_df[param], plot_df['metric'], alpha=0.7)
                    ax.set_xlabel(param)
                    ax.set_ylabel(metric)
                    ax.set_title(f"{param} vs {metric}")
                    
                    # 添加拟合曲线
                    try:
                        # 如果参数是数值型的，尝试拟合曲线
                        if np.issubdtype(plot_df[param].dtype, np.number):
                            from scipy.stats import linregress
                            slope, intercept, r_value, p_value, std_err = linregress(
                                plot_df[param], plot_df['metric']
                            )
                            x = np.array([plot_df[param].min(), plot_df[param].max()])
                            y = intercept + slope * x
                            ax.plot(x, y, 'r--', label=f'R² = {r_value**2:.2f}')
                            ax.legend()
                    except:
                        pass
                    
                    plt.grid(True)
                    plt.tight_layout()
                    plot_path = os.path.join(self.output_dir, f"optimization_{param}_{timestamp}.png")
                    plt.savefig(plot_path, dpi=150)
                    plt.close()
                
                # 创建3D图表(如果有至少2个参数)
                if n_params >= 2:
                    from mpl_toolkits.mplot3d import Axes3D
                    
                    # 选择前两个参数
                    param1, param2 = params[:2]
                    
                    fig = plt.figure(figsize=(10, 8))
                    ax = fig.add_subplot(111, projection='3d')
                    
                    scatter = ax.scatter(
                        plot_df[param1], 
                        plot_df[param2], 
                        plot_df['metric'],
                        c=plot_df['metric'],
                        cmap='viridis',
                        s=50,
                        alpha=0.7
                    )
                    
                    ax.set_xlabel(param1)
                    ax.set_ylabel(param2)
                    ax.set_zlabel(metric)
                    
                    # 标记最佳点
                    best_point = plot_df.iloc[0]
                    ax.scatter(
                        [best_point[param1]],
                        [best_point[param2]],
                        [best_point['metric']],
                        color='red',
                        s=100,
                        marker='*',
                        label='最佳点'
                    )
                    
                    ax.legend()
                    plt.colorbar(scatter)
                    plt.title(f"参数优化 3D 图表: {param1} vs {param2} vs {metric}")
                    
                    plot_path = os.path.join(self.output_dir, f"optimization_3d_{timestamp}.png")
                    plt.savefig(plot_path, dpi=150)
                    plt.close()
                
            except Exception as e:
                logging.error(f"生成优化可视化图表失败: {str(e)}", exc_info=True)
        
        # 创建优化结果报告
        report_path = os.path.join(self.output_dir, f"optimization_report_{timestamp}.html")
        try:
            self._generate_html_report(df_results, metric, higher_is_better, report_path)
        except Exception as e:
            logging.error(f"生成HTML报告失败: {str(e)}", exc_info=True)
        
        logging.info(f"优化报告已保存至: {report_path}")
        return report_path
    
    def _generate_html_report(self, df_results: pd.DataFrame, metric: str,
                            higher_is_better: bool, report_path: str) -> None:
        """生成HTML优化报告"""
        # 导入base64和BytesIO用于图表嵌入
        import base64
        from io import BytesIO
        
        # 排序结果
        if higher_is_better:
            df_sorted = df_results.sort_values('metric', ascending=False)
        else:
            df_sorted = df_results.sort_values('metric', ascending=True)
        
        # 创建可视化图表
        figures = []
        
        try:
            # 只显示前100个结果
            top_results = df_sorted.head(100)
            
            # 对每个参数创建与指标的关系图
            params = [col for col in top_results.columns if col != 'metric']
            for param in params:
                if np.issubdtype(top_results[param].dtype, np.number):
                    fig, ax = plt.subplots(figsize=(8, 5))
                    ax.scatter(top_results[param], top_results['metric'], alpha=0.7)
                    ax.set_xlabel(param)
                    ax.set_ylabel(metric)
                    ax.set_title(f"{param} vs {metric}")
                    ax.grid(True)
                    
                    # 尝试添加拟合曲线
                    try:
                        from scipy.stats import linregress
                        slope, intercept, r_value, p_value, std_err = linregress(
                            top_results[param], top_results['metric']
                        )
                        x = np.array([top_results[param].min(), top_results[param].max()])
                        y = intercept + slope * x
                        ax.plot(x, y, 'r--', label=f'R² = {r_value**2:.2f}')
                        ax.legend()
                    except:
                        pass
                    
                    plt.tight_layout()
                    figures.append((fig, f"{param}_vs_{metric}"))
        except Exception as e:
            logging.error(f"生成优化图表时出错: {str(e)}", exc_info=True)
        
        # 开始构建HTML报告
        html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <meta charset="utf-8">
            <title>参数优化报告</title>
            <style>
                body {{
                    font-family: Arial, sans-serif;
                    margin: 0;
                    padding: 20px;
                    background-color: #f5f5f5;
                }}
                .container {{
                    max-width: 1200px;
                    margin: 0 auto;
                    background-color: white;
                    padding: 20px;
                    box-shadow: 0 0 10px rgba(0,0,0,0.1);
                }}
                h1, h2, h3 {{
                    color: #333;
                }}
                table {{
                    border-collapse: collapse;
                    width: 100%;
                    margin: 20px 0;
                }}
                th, td {{
                    text-align: left;
                    padding: 8px;
                    border-bottom: 1px solid #ddd;
                }}
                th {{
                    background-color: #f2f2f2;
                }}
                tr:hover {{
                    background-color: #f5f5f5;
                }}
                .figure {{
                    margin: 20px 0;
                    text-align: center;
                }}
                .figure img {{
                    max-width: 100%;
                }}
                .highlight {{
                    background-color: #e6f7e6;
                    font-weight: bold;
                }}
                .summary {{
                    margin: 20px 0;
                    padding: 15px;
                    background-color: #f9f9f9;
                    border-left: 5px solid #4CAF50;
                }}
            </style>
        </head>
        <body>
            <div class="container">
                <h1>参数优化报告</h1>
                <p>生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}</p>
                
                <div class="summary">
                    <h2>优化摘要</h2>
                    <p>优化目标: {metric} ({'越高越好' if higher_is_better else '越低越好'})</p>
                    <p>参数数量: {len(params)}</p>
                    <p>评估组合数: {len(df_results)}</p>
                    <p>最佳参数组合:</p>
                    <ul>
        """
        
        # 添加最佳参数
        best_params = df_sorted.iloc[0].to_dict()
        for param, value in best_params.items():
            if param != 'metric':
                html += f"<li>{param}: {value}</li>\n"
        
        html += f"""
                    </ul>
                    <p>最佳{metric}: {best_params['metric']:.6f}</p>
                </div>
                
                <h2>参数优化可视化</h2>
        """
        
        # 添加图表
        for fig, title in figures:
            buf = BytesIO()
            fig.savefig(buf, format='png', dpi=100)
            buf.seek(0)
            img_str = base64.b64encode(buf.read()).decode('utf-8')
            plt.close(fig)
            
            html += f"""
                <div class="figure">
                    <h3>{title}</h3>
                    <img src="data:image/png;base64,{img_str}">
                </div>
            """
        
        # 添加结果表格
        html += """
                <h2>优化结果</h2>
                <p>显示前50个最佳结果:</p>
                <table>
                    <tr>
        """
        
        # 添加表头
        for col in df_sorted.columns:
            html += f"<th>{col}</th>\n"
        
        html += "</tr>\n"
        
        # 添加数据行
        for i, row in df_sorted.iloc[:50].iterrows():
            if i == 0:
                html += '<tr class="highlight">\n'
            else:
                html += "<tr>\n"
                
            for col in df_sorted.columns:
                if col == 'metric':
                    html += f"<td>{row[col]:.6f}</td>\n"
                else:
                    html += f"<td>{row[col]}</td>\n"
                
            html += "</tr>\n"
        
        # 完成HTML
        html += """
                </table>
                
                <h2>总结</h2>
                <p>
                    以上是参数优化的结果。您可以使用最佳参数组合来运行您的策略，
                    也可以根据可视化结果和表格数据进一步调整参数。
                </p>
                <p>
                    注意，最佳参数在此回测期间表现最好，但不一定在未来的市场中同样有效。
                    建议在不同的市场条件下验证策略的稳健性。
                </p>
                
                <div style="margin-top:50px; text-align:center; color:#888;">
                    <p>由QTorch量化交易框架生成</p>
                </div>
            </div>
        </body>
        </html>
        """
        
        # 写入HTML文件
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(html)
    
    def get_best_strategy(self) -> IStrategy:
        """
        获取使用最佳参数的策略实例
        
        返回:
            IStrategy: 策略实例
        """
        if not self.best_params:
            raise ValueError("还没有执行优化或没有找到最佳参数")
        
        return self.strategy_class(**self.best_params)